import random
from typing import List

from agent.agent import Agent
from entity.processor import Processor
from entity.task import Task
from environment.env import Environment
from utils.compute import analyze_schedule_result, calc_execute_time
from utils.gen import gen_test_cases
from utils.plot import plot_gantt


def random_schedule(tasks: List[Task], processors: List[Processor]) -> dict:
    s_map = dict()
    for t in tasks:
        tid = t.id
        r_idx = random.randint(0, len(p_map) - 1)
        pid = processors[r_idx].id
        s_map[tid] = pid
    return s_map


def greedy_schedule(tasks: List[Task], p_map: dict) -> dict:
    s_map = dict()
    m_map = dict()
    for pid in p_map:
        m_map[pid] = 0
    for t in tasks:
        tid = t.id
        min_pid = 0
        min_tot = float('inf')
        for pid in p_map:
            p = p_map[pid]
            exec_time = calc_execute_time(t, p)
            m = m_map[pid]
            tot = m + exec_time
            if tot < min_tot:
                min_tot = tot
                min_pid = pid

        m_map[min_pid] = min_tot
        s_map[tid] = min_pid
    return s_map


def lpt_schedule(tasks: List[Task], p_map: dict) -> dict:
    return greedy_schedule(tasks, p_map)


def spt_schedule(tasks: List[Task], p_map: dict) -> dict:
    return greedy_schedule(tasks, p_map)


def de_schedule(tasks: List[Task], p_map: dict) -> dict:
    return greedy_schedule(tasks, p_map)


def analyze(tasks: List[Task], s_map: dict, p_map: dict, title: str, out_name: str):
    plot_gantt(tasks, s_map, p_map, title, out_name)
    analysis = analyze_schedule_result(tasks, s_map, p_map)
    print(f'[{title}] mean = {analysis.mean}, variance = {analysis.variance}, timeout_cnt = {analysis.timeout_cnt}')


if __name__ == '__main__':
    tasks, processors = gen_test_cases()

    p_map = dict()
    for processor in processors:
        p_map[processor.id] = processor

    rl_tasks = tasks.copy()
    env = Environment(rl_tasks, processors)
    agent = Agent(env)
    for i in range(50):
        agent.play_one_episode(verbose=False)

    s_map = env.schedule_map
    print(s_map)
    analyze(rl_tasks, s_map, p_map, 'RL Schedule Gantt', 'rl')

    rand_tasks = tasks.copy()
    random_s_map = random_schedule(rand_tasks, processors)
    analyze(rand_tasks, random_s_map, p_map, 'Random Schedule Gantt', 'random')

    greedy_tasks = tasks.copy()
    greedy_s_map = greedy_schedule(greedy_tasks, p_map)
    analyze(greedy_tasks, greedy_s_map, p_map,
            'Greedy Schedule Gantt', 'greedy')

    lpt_tasks = tasks.copy()
    lpt_tasks.sort(key=lambda t: t.length(), reverse=True)
    lpt_s_map = lpt_schedule(lpt_tasks, p_map)
    analyze(lpt_tasks, lpt_s_map, p_map, 'LPT Schedule Gantt', 'lpt')

    spt_tasks = tasks.copy()
    spt_tasks.sort(key=lambda t: t.length())
    spt_s_map = spt_schedule(spt_tasks, p_map)
    analyze(spt_tasks, spt_s_map, p_map, 'SPT Schedule Gantt', 'spt')

    de_tasks = tasks.copy()
    de_tasks.sort(key=lambda t: (t.deadline, t.length()))
    de_s_map = de_schedule(de_tasks, p_map)
    analyze(de_tasks, de_s_map, p_map, 'DE Schedule Gantt', 'de')
